/*	
 * Authors: Ken YANG <spng.yang@gmail.com>
 * 	
 * Mammoth client main program
 *
 * Copyright (C) 2007-2008 Ken YANG
 * This program is free software; you can redistribute it and/or 
 * modify it under the terms of the GNU General Public License 
 * as published by the Free Software Foundation, version 2. 
 * THIS SOFTWARE COMES "AS IS" WITHOUT ANY WARRANTIES.
 */

#include "net_basic.h"


/* Wrapper of socket, including error handle */
int Socket(int family, int type, int protocol)
{
	int fd;

	if((fd=socket(family, type, protocol))<0)
		err_sys("socket error");

	return fd;
}


/* Same with Socket */
void Connect(int sockfd, const SA serv_addr, socklen_t addrlen)
{
	if(connect(sockfd, serv_addr, addrlen)<0)
		err_sys("connect error");
}

void Bind(int sockfd, const SA my_addr, socklen_t addrlen)
{
	if(bind(sockfd, my_addr, addrlen)<0)
		err_sys("bind error");
}


void Listen(int sockfd, int backlog)
{
	char *backlog_env;

	/* Allow env variable to overwrite backlog */
	if((backlog_env=getenv("LISTENQ"))!=NULL)
		backlog=atoi(backlog_env);

	if(listen(sockfd, backlog)<0)
		err_sys("listen error");
}

int Accept(int sockfd, SA addr, socklen_t *addrlen)
{
	int conn_fd;

	/* 
	 * If client send RST, between connect(client) 
	 * and accpet(server), then accept set errno 
	 * to "EPROTO" or "ECONNABORTED"(depending on 
	 * the system), we should call accept and block 
	 * to wait for other connections.
	 * */
again:
	if((conn_fd=accept(sockfd, addr, addrlen))<0){
#ifdef EPROTO 
		if(errno == EPROTO || errno ==  ECONNABORTED)
#else
		if(errno == ECONNABORTED)
#endif
			goto again;
		else
			err_sys("accept error");
	}
	return conn_fd;
}


void Close(int sockfd)
{
	if(close(sockfd)<0)
		err_sys("close error");
}

void Getsockname(int sockfd, SA localaddr, socklen_t *addrlen)
{
	if(getsockname(sockfd, localaddr, addrlen)<0)
		err_sys("getsockname error");
}

void Getpeername(int sockfd, SA peeraddr, socklen_t *addrlen)
{
	if(getpeername(sockfd, peeraddr, addrlen)<0)
		err_sys("getpeername error");
}

void Inet_pton(int family, const char *strptr, void *addrptr)
{
	int rc;
	if((rc=inet_pton(family, strptr, addrptr))<0)
		err_sys("inet_pton error for %s", strptr);
	else if(rc==0)
		err_msg("inet_pton error: invalide network address for %s", 
				strptr);
}

void 
Shutdown(int sockfd, int howto)
{
	if(shutdown(sockfd, howto)==-1)
		err_sys("shutdown error");
}

ssize_t Recvfrom(int sockfd, void *buf, size_t nbytes, int flags,
		SA from, socklen_t *addrlen)
{
	int nbyte;
	if((nbyte=recvfrom(sockfd,buf,nbytes,flags,from,addrlen))<0)
		err_sys("recvfrom error");
	return nbyte;
}

ssize_t Sendto(int sockfd, void *buf, size_t nbytes, int flags,
		const SA to, socklen_t addrlen)
{
	int nbyte;
	if((nbyte=sendto(sockfd,buf,nbytes,flags,to,addrlen))<0)
		err_sys("sendto error");
	return nbyte;
}

void Setsockopt(int sockfd, int level, int optname, void *optval,
		socklen_t *optlen)
{
	if(setsockopt(sockfd,level,optname,optval,optlen)==-1)
		err_sys("setsockopt error");
}

void Getsockopt(int sockfd, int level, int optname, void *optval,
		socklen_t *optlen)
{
	if(getsockopt(sockfd,level,optname,optval,optlen)==-1)
		err_sys("getsockopt error");
}



/* Get the protocol family of a socket */
int sockfd_to_family(int sockfd)
{
	struct sockaddr_storage sock;
	socklen_t len=sizeof(sock);

	if(getsockname(sockfd, (SA) &sock, &len)<0)
		return -1;
	return(sock.ss_family);
}


/* Above's wrapper */
int Sockfd_to_family(int sockfd)
{
	int rc;
	if((rc=sockfd_to_family(sockfd))<0)
		err_sys("sockfd_to_family error");
	return rc;
}


/* create tcp socket and connect to a server */
int tcp_connect(const char *hostname, const char *service)
{
	struct addrinfo hints, *result, *index;
	int rc, sockfd;

	bzero(&hints, sizeof(hints));
	hints.ai_socktype=SOCK_STREAM;
	hints.ai_family=AF_UNSPEC;

	if((rc=getaddrinfo(hostname,service,&hints,&result))!=0)
		err_quit(gai_strerror(rc));

	index=result;
	do{
		/* socket error here isn't fatal */
		sockfd=socket(index->ai_family,index->ai_socktype,
				index->ai_protocol);
		if(connect(sockfd, index->ai_addr,index->ai_addrlen)==0)
			break;
		Close(sockfd);	/* ignore failing socket */
	}while((index=index->ai_next)!=NULL);

	/* errno set by final connect */
	if(index==NULL)
		err_sys("tcp_connect error for %s, %s", hostname,
				service);
	freeaddrinfo(result);
	return(sockfd);
}


/* 
 * create tcp socket, bind and listen 
 * "addrlenp" parameter can be NULL, which mean we don't need 
 * the sockaddr size, if not NULL, we use "addrlenp" return the 
 * size of sockaddr
 * */
int tcp_listen(const char *hostname, const char *service, 
		socklen_t *addrlenp)
{
	struct addrinfo hint, *result, *index;
	int rc, listenfd;

	bzero(&hint, sizeof(hint));
	hint.ai_family=AF_UNSPEC;
	hint.ai_socktype=SOCK_STREAM;
	hint.ai_flags=AI_PASSIVE;

	if((rc=getaddrinfo(hostname,service,&hint,&result))!=0)
		err_quit(gai_strerror(rc));

	index=result;
	do{
		listenfd=socket(index->ai_family,index->ai_socktype,
				index->ai_protocol);
		if(listenfd==-1) continue;	/* error, next one */
		Setsockopt(listenfd,SOL_SOCKET,SO_REUSEADDR,1,sizeof(int));
		if(bind(listenfd,index-ai_addr,index->ai_addrlen)==0)
			break;	/* success */
		Close(listenfd);
	}while((index=index->ai_next)!=NULL);

	if(index==NULL)
		err_quit("tcp_listen error for %s\n",hostname);
	
	Listen(listenfd,LISTENQ);

	if(addrlenp)
		*addrlenp=index->ai_addrlen;

	freeaddrinfo(result);
	return(listenfd);
}

/* 
 * create unconnected udp socket 
 * 
 * NOTE:
 * Allocate space for "sa", caller must free this space
 *
 * */
int udp_unconnect(const char *host, const char *service,
		struct sockaddr **sa, socklen_t *lenp)
{
	struct addrinfo hint, *result, *index;
	int rc, sockfd;

	bzero(&hint, sizeof(hint));
	hint.ai_family=AF_UNSPEC;
	hint.ai_socktype=SOCK_DGRAM;

	if((rc=getaddrinfo(host,service,&hint,&result))!=0)
		err_quit(gai_strerror(rc));

	index=result;
	do{
		sockfd=socket(index->ai_family,index->ai_socktype,
				index->ai_protocol);
		if(sockfd>=0) break;
		Close(sockfd);
	}while((index=index->ai_next)!=NULL);

	if(index==NULL)
		err_quit("tcp_listen error for %s\n",host);

	*sa=(SA)Malloc(index->ai_addrlen);
	memcpy(*sa, index->ai_addr, index->ai_addrlen);
	*lenp=index->ai_addrlen;

	freeaddrinfo(result);
	return(sockfd);
}

/* create connected udp socket */
int udp_connect(const char *host, const char *service)
{
	struct addrinfo hint, *result, *index;
	int rc, sockfd;

	bzero(&hint, sizeof(hint));
	hint.ai_family=AF_UNSPEC;
	hint.ai_socktype=SOCK_DGRAM;

	if((rc=getaddrinfo(host,service,&hint,&result))!=0)
		err_quit(gai_strerror(rc));

	index=result;
	do{
		sockfd=socket(index->ai_family,index->ai_socktype,
				index->ai_protocol);
		if(sockfd<0) continue;
		
		if(connect(sockfd,index->ai_addr,index->ai_addrlen)==0)
			break;	/* success */
		Close(sockfd);
	}while((index=index->ai_next)!=NULL);

	if(index==NULL)
		err_quit("tcp_listen error for %s\n",host);

	freeaddrinfo(result);
	return(sockfd);
}

/* create udp socket for server, including bind */
int
udp_server(const char *host, const char *serv, socklen_t *addrlenp)
{
	int sockfd, n;
	struct addrinfo	hints, *res, *ressave;

	bzero(&hints, sizeof(struct addrinfo));
	hints.ai_flags = AI_PASSIVE;
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_DGRAM;

	if ((n = getaddrinfo(host, serv, &hints, &res)) != 0)
		err_quit("udp_server error for %s, %s: %s",
				 host, serv, gai_strerror(n));
	ressave = res;

	do{
		sockfd = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
		if (sockfd < 0)
			continue;	/* error - try next one */

		if (bind(sockfd, res->ai_addr, res->ai_addrlen) == 0)
			break;		/* success */

		Close(sockfd);		/* bind error - close and try next one */
	} while ( (res = res->ai_next) != NULL);

	if (res == NULL)	/* errno from final socket() or bind() */
		err_sys("udp_server error for %s, %s", host, serv);

	if (addrlenp)
		*addrlenp = res->ai_addrlen;	/* return size of protocol address */

	freeaddrinfo(ressave);
	return(sockfd);
}

/* 
 * connect with timeout, for the timeout of 
 * tcp is too long 
 *
 * Parameter:
 *
 * first three: same with "real" connect
 * last one: the seconds of timeout
 * */
int 
connect_timeout(int sockfd, const SA sa, socklen_t salen,
		int nsec)
{
	Sigfunc *oldfunc;
	unsigned int remain_sec;
	int rc;

	oldfunc=Signal(SIGALRM, connect_alarm);

	if((remain_sec=alarm(nsec))>0)
		err_msg("connect_timeout: alarm was already set");

	if((rc=connect(sockfd, sa, salen))==-1){
		Close(sockfd);
		if(errno==EINTR)
			errno=ETIMEDOUT;
	}

	alarm(0);	/* close the alarm */
	if(oldfunc)
		Signal(SIGALRM, oldfunc);
	return rc;
}

static void 
connect_alarm(int signo)
{
	return;
}


/* 
 * connect nonblock version
 *
 * Parameter:
 *
 * first three: same with "real" connect
 * last one: the seconds waiting for finishing connection
 * */
int connect_nonb(int sockfd, const SA sa, socklen_t salen, int nsec)
{
	int flags, rc, error, len;
	fd_set rset, wset;
	struct timeval tval;
	
	flags=Fcntl(sockfd, FD_GETFL, 0);
	Fcntl(sockfd, FD_SETFL, flags | O_NONBLOCK);

	error=0;
	if((rc=connect(sockfd, sa, salen))<0){
		if(errno!=EINPROGRESS)
			return -1;
	}

	/* Do whatever we want, while the connection is taking place */

	/* 
	 * when client and server are in same machine, the connection 
	 * can be established right now.
	 * */
	if(rc==0)
		goto done;

	FD_ZERO(&rset);
	FD_SET(sockfd, &rset);
	wset=rset;
	tval.tv_sec=nsec;
	tval.tv_usec=0;

	/* 
	 * value of the 4th parameter is different between 0 and NULL 
	 * */
	rc=Select(sockfd+1, &rset, &wset, NULL, nsec?&tval:NULL);

	if(rc==0){	/* timeout */
		errno=ETIMEDOUT;
		Close(sockfd);
		return -1;
	}

	if(FD_ISSET(sockfd, &rset) || FD_ISSET(sockfd, &wset)){
		len=sizeof(error);
		if(getsockopt(sockfd,SOL_SOCKET,SO_ERROR,&error,*len)<0)
			/* 
			 * consider about portable, solaris and BSD are 
			 * different in this, which we handle here. 
			 * Either solaris and bsd, we all return error 
			 * in "errno", and retur value as -1
			 * */
			return -1;
	}else
		err_quit("select error: sockfd not set");

done:
	Fcntl(sockfd, FD_SETFL, flags);
	if(error){
		Close(sockfd);
		errno=error;
		return -1;
	}

	return 0;
}


/*
 * Convert hostname or service to socket address structure.
 *
 * NOTE:
 * Since getaddrinfo don't free the result addrinfo, the 
 * caller must free it by himself, through freeaddrinfo()
 * */

struct addrinfo *
Host_serv(const char *host, const char *serv, int family, int socktype)
{
	int rc;
	struct addrinfo hint, *result;

	bzero(&hint,sizeof(struct addrinfo));
	hint.ai_flags=AI_CANONNAME;	/* always return canonical name */
	hint.ai_family=family;          /* 0, AF_INET, AF_INET6, etc. */
	hint.ai_socktype = socktype;   /* 0, SOCK_STREAM, SOCK_DGRAM, etc. */

	if((rc=getaddrinfo(host,serv,&hint,&result))!=0)
		err_quit("host_serv error for %s, %s: %s",
				(host == NULL) ? "(no hostname)" : host,
				(serv == NULL) ? "(no service name)" : serv,
				gai_strerror(n));
	return(res);	/* return the first on the linked list */
}


/* 
 * Because inet_ntop must speicific address of socket structure, 
 * we have to call inet_ntop with different options for various 
 * address family. 
 *
 * So we define sock_ntop(), sock_ntop_host(), and etc to avoid 
 * this problem. sock_ntop()'s conversion include port, but 
 * sock_ntop_host() not.
 *
 * FIXME: sock_ntop() is not thread-safe, because the return 
 * value point to a static memory
 * */
char *
sock_ntop(const struct sockaddr *sa, socklen_t salen)
{
	char portstr[8];
	static char str[128];	/* Unix domain is largest */

	switch (sa->sa_family) {
	case AF_INET: {
		struct sockaddr_in	*sin = (struct sockaddr_in *) sa;

		if (inet_ntop(AF_INET, &sin->sin_addr, str, sizeof(str)) == NULL)
			return(NULL);
		if (ntohs(sin->sin_port) != 0) {
			snprintf(portstr, sizeof(portstr), ":%d", ntohs(sin->sin_port));
			strcat(str, portstr);
		}
		return(str);
	}
/* end sock_ntop */

#ifdef	IPV6
	case AF_INET6: {
		struct sockaddr_in6	*sin6 = (struct sockaddr_in6 *) sa;

		str[0] = '[';
		if (inet_ntop(AF_INET6, &sin6->sin6_addr, str + 1, sizeof(str) - 1) == NULL)
			return(NULL);
		if (ntohs(sin6->sin6_port) != 0) {
			snprintf(portstr, sizeof(portstr), "]:%d", ntohs(sin6->sin6_port));
			strcat(str, portstr);
			return(str);
		}
		return (str + 1);
	}
#endif

#ifdef	AF_UNIX
	case AF_UNIX: {
		struct sockaddr_un	*unp = (struct sockaddr_un *) sa;

			/* OK to have no pathname bound to the socket: happens on
			   every connect() unless client calls bind() first. */
		if (unp->sun_path[0] == 0)
			strcpy(str, "(no pathname bound)");
		else
			snprintf(str, sizeof(str), "%s", unp->sun_path);
		return(str);
	}
#endif

#ifdef	HAVE_SOCKADDR_DL_STRUCT
	case AF_LINK: {
		struct sockaddr_dl	*sdl = (struct sockaddr_dl *) sa;

		if (sdl->sdl_nlen > 0)
			snprintf(str, sizeof(str), "%*s (index %d)",
					 sdl->sdl_nlen, &sdl->sdl_data[0], sdl->sdl_index);
		else
			snprintf(str, sizeof(str), "AF_LINK, index=%d", sdl->sdl_index);
		return(str);
	}
#endif
	default:
		snprintf(str, sizeof(str), "sock_ntop: unknown AF_xxx: %d, len %d",
				 sa->sa_family, salen);
		return(str);
	}
    return (NULL);
}

char *
Sock_ntop(const struct sockaddr *sa, socklen_t salen)
{
	char	*ptr;

	if ( (ptr = sock_ntop(sa, salen)) == NULL)
		err_sys("sock_ntop error");	/* inet_ntop() sets errno */
	return(ptr);
}


/* 
 * convert the "host part of sockaddr to presentation, without 
 * port.
 * */
char *
sock_ntop_host(const struct sockaddr *sa, socklen_t salen)
{
    static char str[128];		/* Unix domain is largest */

	switch (sa->sa_family) {
	case AF_INET: {
		struct sockaddr_in *sin = (struct sockaddr_in *) sa;

		if (inet_ntop(AF_INET, &sin->sin_addr, str, 
					sizeof(str)) == NULL)
			return(NULL);
		return(str);
	}

#ifdef	IPV6
	case AF_INET6: {
		struct sockaddr_in6	*sin6 = (struct sockaddr_in6 *) sa;

		if (inet_ntop(AF_INET6, &sin6->sin6_addr, str, sizeof(str)) == NULL)
			return(NULL);
		return(str);
	}
#endif

#ifdef	AF_UNIX
	case AF_UNIX: {
		struct sockaddr_un	*unp = (struct sockaddr_un *) sa;

			/* OK to have no pathname bound to the socket: happens on
			   every connect() unless client calls bind() first. */
		if (unp->sun_path[0] == 0)
			strcpy(str, "(no pathname bound)");
		else
			snprintf(str, sizeof(str), "%s", unp->sun_path);
		return(str);
	}
#endif

#ifdef	HAVE_SOCKADDR_DL_STRUCT
	case AF_LINK: {
		struct sockaddr_dl	*sdl = (struct sockaddr_dl *) sa;

		if (sdl->sdl_nlen > 0)
			snprintf(str, sizeof(str), "%*s",
					 sdl->sdl_nlen, &sdl->sdl_data[0]);
		else
			snprintf(str, sizeof(str), "AF_LINK, index=%d", sdl->sdl_index);
		return(str);
	}
#endif
	default:
		snprintf(str, sizeof(str), "sock_ntop_host: unknown AF_xxx: %d, len %d",
				 sa->sa_family, salen);
		return(str);
	}
    return (NULL);
}

char *
Sock_ntop_host(const struct sockaddr *sa, socklen_t salen)
{
	char	*ptr;

	if ( (ptr = sock_ntop_host(sa, salen)) == NULL)
		err_sys("sock_ntop_host error");	/* inet_ntop() sets errno */
	return(ptr);
}

/* as function name said */
int
sock_get_port(const struct sockaddr *sa, socklen_t salen)
{
	switch (sa->sa_family) {
	case AF_INET: {
		struct sockaddr_in	*sin = (struct sockaddr_in *) sa;

		return(sin->sin_port);
	}

#ifdef	IPV6
	case AF_INET6: {
		struct sockaddr_in6	*sin6 = (struct sockaddr_in6 *) sa;

		return(sin6->sin6_port);
	}
#endif
	}

    return(-1);		/* ??? */
}

